#include "MyImage.hpp"
#include <cassert>
#include <fstream>
#include <vector>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <cstdint>

Image::Image(int w, int h, int c, ImageFormat f)
    : width(w), height(h), channels(c), format(f), ownsMemory(true)
{
    if (f == ImageFormat::LDR)
    {
        ldrBuffer = new unsigned char[w * h * c];
        hdrBuffer = nullptr;
    }
    else
    {
        hdrBuffer = new float[w * h * c];
        ldrBuffer = nullptr;
    }
}

Image::~Image()
{
    if (ownsMemory)
    {
        delete[] ldrBuffer;
        delete[] hdrBuffer;
    }

    ldrBuffer = nullptr;
    hdrBuffer = nullptr;
}

std::string Image::GetExtension(const std::string& filename)
{
    size_t p = filename.find_last_of('.');
    if (p == std::string::npos) return "";
    std::string ext = filename.substr(p + 1);
    for (auto& c : ext) c = tolower(c);
    return ext;
}

Image* Image::Load(const std::string& filename)
{
    std::string ext = GetExtension(filename);

    if (ext == "tga") return LoadTGA(filename);
    if (ext == "hdr") return LoadHDR(filename);

    assert(!"Unsupported image format");
    return nullptr;
}

void Image::Save(const std::string& filename)
{
    std::string ext = GetExtension(filename);

    if (ext == "tga") { SaveTGA(filename); return; }
    if (ext == "hdr") { SaveHDR(filename); return; }

    assert(!"Unsupported image format");
}

void Image::FlipH()
{
    int halfWidth = width / 2;
    for (int row = 0; row < height; ++row)
    {
        for (int col = 0; col < halfWidth; ++col)
        {
            int oppositeCol = width - 1 - col;

            if (format == ImageFormat::LDR)
            {
                unsigned char* p1 = GetLdrPixel(row, col);
                unsigned char* p2 = GetLdrPixel(row, oppositeCol);
                for (int k = 0; k < channels; ++k)
                    std::swap(p1[k], p2[k]);
            }
            else 
            {
                float* p1 = GetHdrPixel(row, col);
                float* p2 = GetHdrPixel(row, oppositeCol);
                for (int k = 0; k < channels; ++k)
                    std::swap(p1[k], p2[k]);
            }
        }
    }
}

void Image::FlipV()
{
    int halfHeight = height / 2;
    for (int row = 0; row < halfHeight; ++row)
    {
        int oppositeRow = height - 1 - row;
        for (int col = 0; col < width; ++col)
        {
            if (format == ImageFormat::LDR)
            {
                unsigned char* p1 = GetLdrPixel(row, col);
                unsigned char* p2 = GetLdrPixel(oppositeRow, col);
                for (int k = 0; k < channels; ++k)
                    std::swap(p1[k], p2[k]);
            }
            else 
            {
                float* p1 = GetHdrPixel(row, col);
                float* p2 = GetHdrPixel(oppositeRow, col);
                for (int k = 0; k < channels; ++k)
                    std::swap(p1[k], p2[k]);
            }
        }
    }
}

unsigned char* Image::GetLdrPixel(int row, int col)
{
    return &ldrBuffer[(row * width + col) * channels];
}

float* Image::GetHdrPixel(int row, int col)
{
    return &hdrBuffer[(row * width + col) * channels];
}

void Image::SetPixel(int row, int col, const float* value)
{
    assert(format == ImageFormat::HDR);
    assert(row >= 0 && row < height);
    assert(col >= 0 && col < width);

    float* dst = GetHdrPixel(row, col);
    for (int k = 0; k < channels; ++k)
    {
        dst[k] = value[k];
    }
}

void Image::SetPixel(int row, int col, const unsigned char* value)
{
    assert(format == ImageFormat::LDR);
    assert(row >= 0 && row < height);
    assert(col >= 0 && col < width);

    unsigned char* dst = GetLdrPixel(row, col);
    for (int k = 0; k < channels; ++k)
    {
        dst[k] = value[k];
    }
}

struct TGAHeader
{
    uint8_t  idLength;
    uint8_t  colorMapType;
    uint8_t  dataType;     
    uint16_t colorMapStart;
    uint16_t colorMapLength;
    uint8_t  colorMapDepth;
    uint16_t xOrigin;
    uint16_t yOrigin;
    uint16_t width;
    uint16_t height;
    uint8_t  bitsPerPixel;
    uint8_t  descriptor;
};
#pragma pack(pop)

Image* Image::LoadTGA(const std::string& filename)
{
    std::ifstream f(filename, std::ios::binary);
    if (!f) { assert(!"Cannot open TGA file"); return nullptr; }

    TGAHeader h;
    f.read(reinterpret_cast<char*>(&h), sizeof(h));

    if (h.dataType != 2 && h.dataType != 3)
    {
        assert(!"Only uncompressed RGB/RGBA/Gray TGA supported");
        return nullptr;
    }

    int channels = (h.bitsPerPixel / 8);
    if (!(channels == 1 || channels == 3 || channels == 4))
    {
        assert(!"Unsupported TGA pixel format");
        return nullptr;
    }

    f.seekg(h.idLength, std::ios::cur);

    Image* img = new Image(h.width, h.height, channels, ImageFormat::LDR);

    size_t dataSize = (size_t)h.width * h.height * channels;
    f.read(reinterpret_cast<char*>(img->ldrBuffer), dataSize);

    if (channels >= 3)
    {
        for (int i = 0; i < h.width * h.height; ++i)
        {
            std::swap(img->ldrBuffer[i * channels], img->ldrBuffer[i * channels + 2]);
        }
    }

    return img;
}

void Image::SaveTGA(const std::string& filename)
{
    assert(format == ImageFormat::LDR);

    std::ofstream f(filename, std::ios::binary);
    if (!f) { assert(!"Cannot write TGA"); return; }

    TGAHeader h{};
    h.dataType = (channels == 1 ? 3 : 2);
    h.width = width;
    h.height = height;
    h.bitsPerPixel = channels * 8;

    f.write(reinterpret_cast<char*>(&h), sizeof(h));

    for (int i = 0; i < width * height; ++i)
    {
        unsigned char* px = &ldrBuffer[i * channels];

        if (channels == 1)
        {
            f.write((char*)px, 1);
        }
        else
        {
            unsigned char bgr[4];
            bgr[0] = px[2];
            bgr[1] = px[1];
            bgr[2] = px[0];
            if (channels == 4)
                bgr[3] = px[3];
            f.write((char*)bgr, channels);
        }
    }
}

static void rgbeToFloat(float* dst, const unsigned char rgbe[4])
{
    if (rgbe[3] == 0)
    {
        dst[0] = dst[1] = dst[2] = 0.0f;
        return;
    }

    float f = std::ldexp(1.0f, rgbe[3] - (128 + 8)); // 2^(E-136)

    dst[0] = rgbe[0] * f;
    dst[1] = rgbe[1] * f;
    dst[2] = rgbe[2] * f;
}

Image* Image::LoadHDR(const std::string& filename)
{
    std::ifstream f(filename, std::ios::binary);
    if (!f) { assert(!"Cannot open HDR"); return nullptr; }

    std::string line;
    std::getline(f, line);
    if (line != "#?RADIANCE")
    {
        assert(!"Not a Radiance HDR file");
        return nullptr;
    }

    while (true)
    {
        std::getline(f, line);
        if (line.empty()) break;
    }

    int w, h;
    f >> line >> h >> line >> w;
    f.get(); 

    Image* img = new Image(w, h, 3, ImageFormat::HDR);

    for (int y = 0; y < h; y++)
    {
        for (int x = 0; x < w; x++)
        {
            unsigned char rgbe[4];
            f.read((char*)rgbe, 4);

            float* out = img->GetHdrPixel(y, x);
            rgbeToFloat(out, rgbe);
        }
    }

    return img;
}

static void floatToRgbe(unsigned char rgbe[4], const float* rgb)
{
    float v = std::max({ rgb[0], rgb[1], rgb[2] });

    if (v < 1e-32f)
    {
        rgbe[0] = rgbe[1] = rgbe[2] = rgbe[3] = 0;
        return;
    }

    int e;
    float m = std::frexp(v, &e); 

    m *= 256.0f / v;

    rgbe[0] = rgb[0] * m;
    rgbe[1] = rgb[1] * m;
    rgbe[2] = rgb[2] * m;
    rgbe[3] = e + 128;
}

void Image::SaveHDR(const std::string& filename)
{
    assert(format == ImageFormat::HDR);

    std::ofstream f(filename, std::ios::binary);
    if (!f) { assert(!"Cannot write HDR"); return; }

    f << "#?RADIANCE\nFORMAT=32-bit_rle_rgbe\n\n";
    f << "-Y " << height << " +X " << width << "\n";

    for (int y = 0; y < height; y++)
    {
        for (int x = 0; x < width; x++)
        {
            float* px = GetHdrPixel(y, x);
            unsigned char rgbe[4];
            floatToRgbe(rgbe, px);
            f.write((char*)rgbe, 4);
        }
    }
}
